﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.Xna.Framework;
using System.Diagnostics;

namespace Space
{
    class Node<T> where T : IPosition
    {
        Vector2 location;
        Vector2 size;

        List<T> values;

        Node<T>[] children;

        enum Location
        {
            TopLeft = 0 << 1 | 0,
            TopRight = 0 << 1 | 1,
            BottomLeft = 1 << 1 | 0,
            BottomRight = 1 << 1 | 1
        }

        int level;
        const int MaxLevel = 20;

        public Node(int limit, Vector2 loc, Vector2 size, int level=1)
        {
            values = new List<T>(limit);
            this.size = size;
            location = loc;
            this.level = level;
        }

        public void Add(T position)
        {
            if (values != null && (values.Count < values.Capacity || level == MaxLevel)) values.Add(position);
            else
            {
                Vector2 center = location + size / 2;

                if (children == null)
                {
                    children = new Node<T>[4];
                    children[(int)Location.TopLeft] = new Node<T>(values.Capacity, location, size / 2, level + 1);
                    children[(int)Location.TopRight] = new Node<T>(values.Capacity, location + new Vector2(size.X / 2, 0), size / 2, level + 1);
                    children[(int)Location.BottomLeft] = new Node<T>(values.Capacity, location + new Vector2(0, size.Y / 2), size / 2, level + 1);
                    children[(int)Location.BottomRight] = new Node<T>(values.Capacity, location + size / 2, size / 2, level + 1);
                }

                AddToChild(position);

                if (values != null)
                {
                    foreach (T p in values)
                        AddToChild(p);

                    values = null;
                }
            }
        }

        void AddToChild(T position)
        {
            int insertLocation = 0;
            if (position.Position.X > location.X + size.X / 2) insertLocation |= 1;
            if (position.Position.Z > location.Y + size.Y / 2) insertLocation |= 1 << 1;

            children[insertLocation].Add(position);

            Debug.Assert(position.Position.X >= children[insertLocation].location.X && 
                position.Position.Z >= children[insertLocation].location.Y && 
                position.Position.X <= children[insertLocation].location.X + children[insertLocation].size.X &&
                position.Position.Z <= children[insertLocation].location.Y + children[insertLocation].size.Y);
        }

        public bool Intercects(Vector2 center, float radius)
        {
            Vector2 loc = location - new Vector2(radius);
            Vector2 siz = size + new Vector2(radius * 2);
            return center.X >= loc.X && center.Y >= loc.Y && center.X <= loc.X + siz.X && center.Y <= loc.Y + siz.Y;
        }

        public IEnumerable<IPosition> Find(Vector2 center, float radius)
        {
            if (Intercects(center, radius))
            {
                if (values != null)
                {
                    foreach (IPosition p in values)
                        if ((new Vector2(p.Position.X, p.Position.Z) - center).LengthSquared() < radius * radius)
                            yield return p;
                }
                else
                {
                    foreach (Node<T> n in children)
                        foreach (IPosition v in n.Find(center, radius))
                            yield return v;
                }
            }
        }
    }

    class QuadTree<T> where T : IPosition
    {
        Node<T> root;

        int limit;
        float size;

        public QuadTree(float size)
        {
            this.size = size;
            limit = 10;
            root = new Node<T>(limit, Vector2.Zero - Vector2.One * size, Vector2.One * size * 2);
        }

        public void Reset()
        {
            root = new Node<T>(limit, Vector2.Zero - Vector2.One * size, Vector2.One * size * 2);
        }

        public void Insert(T obj)
        {
            if (obj.Position.X >= -size && obj.Position.Z >= -size && obj.Position.X <= size && obj.Position.Z <= size)
                root.Add(obj);
        }

        public IEnumerable<T> Find(Vector3 center, float radius)
        {
            foreach (T p in root.Find(new Vector2(center.X, center.Z), radius))
                yield return p;
        }

        public IEnumerable<T> Find(Vector3 center, float radius, Predicate<T> predicate)
        {
            foreach (T p in root.Find(new Vector2(center.X, center.Z), radius))
                if (predicate(p))
                    yield return p;
        }
    }
}
